import math

from torch import  nn,Tensor

class Embedding(nn.Module):
    def __init__(self,vocab_size:int, d_model:int) -> None:
        '''

        :param vocab_size: 词库大小
        :param d_model: embedding的向量维度
        '''
        super().__init__()
        self.embed = nn.Embedding(vocab_size,d_model)
        self.sqrt_d_model = math.sqrt(d_model)


    def forward(self,x:Tensor) -> Tensor:
        '''

        :param x: 输入的Tensor,(batch_size,seq_length)
        :return: Tensor，shape为：(batch_size,seq_length,d_model)
        '''

        ## 论文里面，初始化以后要乘以 d_model的平方根
        return self.embed(x) * self.sqrt_d_model